import math
import numpy as np
import matplotlib.pyplot as plt
import csv
#from scipy.stats import poisson

omega = 0.
i = 3.
m = 3
n = 80
h = 1.

M = 20 #Mmax
Mmin = 2

t_f = n * h
tau = m * h

Delta = 0.8
kappa = 5.
sigma = 0.5
nu = 1.
d_avg = 4.
alpha = 0.01
mu = 0.0125
xi = 20.
P_RS = 0.05
P_IR = 0.25
P = 0.2

D = {}
S = {}
E = {}
I = {}
R = {}
Theta = {}
S[0] = 1. - P
E[0] = 0.
I[0] = P
R[0] = 0.
Theta[0] = 0.

Lambda = {}
Eta = {}
Nu_i = {}
Phi = {}


def factorial(n):
    a = 1
    for i in range(1, n + 1):
        a = a * i
    return a

#possion
def p(i):
    return (d_avg ** i) * (math.exp(-1. * d_avg)) / factorial(i)

def delta(i):
    pSum = 0
    for ii in range(Mmin, M + 1):
        pSum += p(ii)
    return p(i) / pSum

def gamma(i):
    return kappa * (i ** sigma) / (1. + nu * (i**sigma))

def Ci():
    Sum = 0
    for ii in range(Mmin, M + 1):
        Sum += delta(ii) * gamma(ii)
    return Sum / d_avg

def chi(t):
    if t >=0 and t <= (t_f - tau):
        return 1.
    else:
        return 0.

def P_SI(i):
    return alpha * i

def CalcTheta(j):
        #Theta[j+1] = Theta[0]
        #return
        Lambda[n - j - 1] = Lambda[n - j] - h * (Lambda[n - j] * (P_SI(i) * Ci() * I[j+1] + mu) \
            - Eta[n-j] * (1. - omega) * P_SI(i) * Ci() * I[j+1] \
            - Nu_i[n-j] * omega * P_SI(i) *Ci() * I[j+1] \
            - chi(n-j) * (Eta[n-j+m]*(omega - 1.)* P_SI(i)*Ci()*I[j+1]*math.exp(-1. * mu *tau) \
                + Nu_i[n-j+m]*(1. - omega)* P_SI(i)*Ci()*I[j+1]*math.exp(-1. * mu * tau)))

        Eta[n - j - 1] = Eta[n - j] - h * (mu * Eta[n-j])

        Nu_i[n - j - 1] = Nu_i[n - j] - h * (-1. + Lambda[n-j] * P_SI(i) * S[j + 1] *Ci() \
            + Eta[n-j] *(1. - omega) * P_SI(i) * S[j+1] *Ci()\
            - Nu_i[n - j] *(mu + Theta[j] - omega * P_SI(i) * S[j+1] * Ci()) \
            + Phi[n-j] * Theta[j] \
            - chi(n-j) * (Eta[n-j+m]*(omega - 1.)*Ci()*S[j+1]*math.exp(-1. * mu *tau) \
                + Nu_i[n-j+m]*(1. - omega)*Ci()*S[j+1]*math.exp(-1. * mu * tau)))

        Phi[n-j-1] = Phi[n-j] - h * (Lambda[n-j] * P_RS + Phi[n-j] * (P_RS + mu))
        
        D[j+1] = (Nu_i[n-j-1] - Phi[n-j-1]) * I[j+1] / xi
        Theta[j+1] = min(max(0., D[j+1]), Delta)

def Step1():
    for j in range(-m,0):
        S[j] = S[0]
        E[j] = E[0]
        I[j] = I[0]
        R[j] = R[0]
        Theta[j] = Theta[0]
    for j in range(n, n+m+1):
        Lambda[j] = 0.
        Eta[j] = 0.
        Nu_i[j] = 0.
        Phi[j] = 0.

def Step2():
    for j in range(0,n):
        S[j + 1] = S[j] + h * (mu + P_RS * R[j] - P_SI(i) * Ci() * I[j] * S[j] - mu * S[j])
        S[j + 1] = min(max(0., S[j + 1]), 1.)

        E[j + 1] = E[j] + h * ((1. - omega) * P_SI(i) * Ci() * I[j] * S[j] \
            - (1. - omega) * P_SI(i) * Ci() * I[j-m] * S[j-m] * math.exp(-1. * mu * tau) - mu * E[j])
        E[j + 1] = min(max(0., E[j+1]), 1.)

        I[j + 1] = I[j] + h * (omega * P_SI(i) * Ci() * I[j] * S[j] - Theta[j] * I[j] \
            + (1. - omega) * P_SI(i) * Ci() * I[j-m] * S[j-m] * math.exp(-1. * mu * tau) - mu * I[j])
        I[j + 1] = min(max(0., I[j + 1]), 1.)

        R[j + 1] = R[j] + h * (Theta[j] * I[j] - P_RS * R[j] - mu * R[j])
        R[j + 1] = min(max(0., R[j + 1]), 1.)

        CalcTheta(j)


def Step3():
    pass

#def test():
#    data_binom = poisson.rvs(mu=4, size =10000)
#    print(data_binom[1,20])

if __name__=='__main__':
    #test()
    Step1()
    Step2()

    def Show(dict, lbl):
        l1 = []
        for i in range(0, n):
            l1.append(dict[i])
        plt.plot(range(0,n), l1, label=lbl)
    
    with open(r'./test___.csv', 'a',newline="") as f:
        writer = csv.writer(f)
        writer.writerow(['S','E','I','R','Theta'])
        for j in range(0,n):
            writer.writerow([S[j], E[j], I[j], R[j],Theta[j]]) 

    plt.figure(1)
    Show(S, 'S')
    Show(E, 'E')
    Show(I, 'I')
    Show(R, 'R')
    
    plt.legend()
    
    plt.figure(2)
    Show(Theta, 'Theta')
    plt.legend()
    plt.show()



